import faiss
import numpy as np
import json
import os
from pathlib import Path
import torch

class FaissDatabase:
    def __init__(self, index_path="faiss_index.bin", meta_path="metadata.json"):
        self.index_path = index_path
        self.meta_path = meta_path
        self.index = None
        self.metadata = []
        self._initialize_index()

    def _initialize_index(self):
        # 尝试加载已有索引
        if Path(self.index_path).exists():
            if faiss.get_num_gpus() > 0:
                self.index = faiss.read_index(self.index_path, faiss.GPU_0)
            else:
                self.index = faiss.read_index(self.index_path)
            with open(self.meta_path, 'r') as f:
                self.metadata = json.load(f)
        else:
            # 创建新索引
            self._create_new_index()

    def _create_new_index(self):
        dim = 512  # ResNet-18特征维度
        if faiss.get_num_gpus() > 0:
            self.index = faiss.index_factory(dim, "Flat", faiss.METRIC_INNER_PRODUCT)
            self.index = faiss.index_cpu_to_all_gpus(self.index)
        else:
            self.index = faiss.IndexFlatIP(dim)

    def add_feature(self, feature, image_path):
        # 标准化特征向量
        faiss.normalize_L2(feature.reshape(1, -1))
        self.index.add(feature.reshape(1, -1))
        self.metadata.append(image_path)
        self._save_to_disk()

    def search(self, query_feature, top_k=10, threshold=0.7):
        faiss.normalize_L2(query_feature.reshape(1, -1))
        distances, indices = self.index.search(query_feature.reshape(1, -1), top_k)
        
        results = []
        for i in range(top_k):
            idx = indices[0][i]
            if idx == -1 or distances[0][i] < threshold:
                continue
            results.append({
                "path": self.metadata[idx],
                "similarity": float(distances[0][i])
            })
        return sorted(results, key=lambda x: x["similarity"], reverse=True)

    def _save_to_disk(self):
        # 保存索引文件
        if torch.cuda.is_available():
            faiss.write_index(faiss.index_gpu_to_cpu(self.index), self.index_path)
        else:
            faiss.write_index(self.index, self.index_path)
        # 保存元数据
        with open(self.meta_path, 'w') as f:
            json.dump(self.metadata, f)

    def get_index_size(self):
        return self.index.ntotal